import torch
from typing import Union

class Params:
    paths = {}
    model:str = None
    dst:str = None
    optimizer:str = None
    E:int = None
    BS:int = None
    device:str = None
    lr:float = None
    decay:float = None
    momentum:float = None
    _additional_apporach_params = {}


    def __init__(self,args,paths) -> None:
        self.paths = paths
        self.model = args.model
        self.dst = args.data_set
        self.optimizer = args.optimizer
        self.E = args.epochs
        self.BS = args.batch_size
        self.lr = args.lr
        self.decay = args.decay
        self.momentum = args.momentum
        self.pretrained = args.pretrained
        self.freeze_ftrs = args.freeze_ftrs
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

    def append_addi_params(self, addi_name:str, addi_params) -> None:
        self._additional_apporach_params[addi_name] = addi_params


    def get_addi_params(self) -> dict:
        return self._additional_apporach_params